"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/6 21:47
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    generate_rgbd_pointcloud.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    generate_rgbd_pointcloud_VER0.1

------------------------------------------------------------------
# @Description:

------------------------------------------------------------------
"""
import os
import cv2
from tqdm import tqdm
from natsort import natsorted
from utils import Utils, EdgeDenoiser
from align import RMSLAlign
from parser_cam_param import ParserCamParam


def get_all_depth_test_image_path(root_dir, key1='_depth', key2='.png', skipPrecision=False):
    if not os.path.exists(root_dir):
        print(' [ERROR] Test directory not exists!')
        return

    # get all path
    img_path_list = Utils().findFilesWithPathlibRecursive(root_dir, key1, key2)
    img_path_list = [img_path for img_path in img_path_list if '_align' not in img_path]

    if skipPrecision:
        # skip precision category
        img_path_list = [img_path for img_path in img_path_list if 'precision' not in img_path]

    return natsorted(img_path_list)


def get_rgb_img_path(env_dp_img_path: str, key='/rk_result_0708'):
    if not os.path.exists(env_dp_img_path):
        print(' [ERROR] Abnormal disp image path.')
        return

    rgb_img_path = env_dp_img_path.replace(key, '').replace('_depth', '_color')

    return rgb_img_path


def generate_depth_images(json_path: str, root_dir: str, dp_w: int, dp_h: int):
    if not os.path.exists(json_path):
        print(' [ERROR] Check json file path.')
        return

    if not os.path.exists(root_dir):
        print(' [ERROR] Test directory is empty')
        return

    # init
    parser = ParserCamParam(json_path)

    # get module param
    cam_param, ratio_w, ratio_h = parser.parser_cam_param(dp_w, dp_h)

    # get all disp image
    img_path_list = get_all_depth_test_image_path(root_dir, key1='_disp', key2='.png', skipPrecision=False)
    depth_img_path_list = []

    # gen depth image
    for idx, disp_path in enumerate(tqdm(img_path_list, position=0, desc=' Gen depth...', colour='cyan')):
        depth_path = disp_path.replace('_disp', '_depth')
        if os.path.exists(depth_path):
            depth_img_path_list.append(depth_path)
            continue

        disp_img = cv2.imread(disp_path, cv2.IMREAD_UNCHANGED)

        dp_img = Utils().disp_to_depth(disp_img,
                                       cam_param['q_coeffs']['Q23'],
                                       cam_param['q_coeffs']['Q32'],
                                       cam_param['q_coeffs']['Q33'],
                                       sub_pixel_value=64,
                                       zoom_ratio=1.0,
                                       max_dis=3000)
        cv2.imwrite(depth_path, dp_img)
        depth_img_path_list.append(depth_path)

    return depth_img_path_list


def generate_all_rgbd_align_pointcloud(json_path: str, root_dir: str, depth_img_path_list: list, dp_w, dp_h, rgb_w, rgb_h, subDir='/rk_result_0719', rotate_angle=0, enReGetList=False):
    if not os.path.exists(json_path):
        print(' [ERROR] Check json file path.')
        return

    if len(depth_img_path_list) <= 0:
        print(' [ERROR] Test directory is empty')
        return

    # init
    parser = ParserCamParam(json_path)
    aligner = RMSLAlign(src_w=dp_w, src_h=dp_h,
                        dst_w=rgb_w, dst_h=rgb_h,
                        extend_w_pixels=256,
                        extend_h_pixels=256,
                        rotate_angle=rotate_angle)

    # get module param
    cam_param, ratio_w, ratio_h = parser.parser_cam_param(dp_w, dp_h)
    aligner.create_lut(cam_param['r_coeffs'])

    if enReGetList:
        depth_img_path_list = get_all_depth_test_image_path(root_dir, key1='_depth', key2='.png', skipPrecision=False)

    for idx, dp_path in enumerate(tqdm(depth_img_path_list, position=0, desc=' Align and Gen PointCloud...', colour='red')):

        rgb_path = get_rgb_img_path(dp_path, key=subDir)
        ply_path = dp_path.replace('.png', '.ply').replace('_depth', '_pointcloud')

        rgb_img = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
        dp_img = cv2.imread(dp_path, cv2.IMREAD_UNCHANGED)

        # 640x400 enable
        if dp_img.shape[0] != dp_h or dp_img.shape[1] != dp_w:
            dp_img = cv2.resize(dp_img, (dp_w, dp_h), interpolation=cv2.INTER_NEAREST)

        # dp_img = Utils().filter_depth_image_edge(dp_img, threshold=200)
        dp_img = EdgeDenoiser().denoiseDepthEdge(dp_img, depth_range=[], threshold_range=[], enMedFilter=False)

        # clear flag
        dp_img[0, :] = 0
        # align depth to rgb
        align_dp_img = aligner.align_dp_to_rgb(dp_img, cam_param['t_coeffs'])
        cv2.imwrite(ply_path.replace('.ply', '_align.png'), align_dp_img)

        # generate color PointCloud
        color_points = Utils().aligned_dp_to_color_point_cloud(rgb_img,
                                                               align_dp_img,
                                                               cam_param['color_intrinsics'],
                                                               depth_ratio=1.0,
                                                               dis_near=100,
                                                               dis_far=3000)
        # Save PointCloud
        Utils().save_ply(ply_path, color_points)


if __name__ == '__main__':
    """--------------------------------------------------------------------------------------------------------"""

    test_env_dir = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/rk_result_0708"
    json_file_path = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/rgb1920x1080_ir1280x800.json"

    """--------------------------------------------------------------------------------------------------------"""
    # Need to correct setting.
    disp_img_w = 1280
    disp_img_h = 800

    dp_image_w = 1280
    dp_image_h = 800
    rgb_image_w = 1920
    rgb_image_h = 1080
    """--------------------------------------------------------------------------------------------------------"""

    dp_img_path_list = generate_depth_images(json_file_path, test_env_dir, disp_img_w, disp_img_h)
    generate_all_rgbd_align_pointcloud(json_file_path,
                                       test_env_dir,
                                       dp_img_path_list,
                                       dp_image_w,
                                       dp_image_h,
                                       rgb_image_w,
                                       rgb_image_h,
                                       subDir='/rk_result_0708',
                                       rotate_angle=0,
                                       enReGetList=False)
